package com.mezz.bones.framework.sse.service;

import org.springframework.stereotype.Service;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

@Service
public class SseSessionManager {
    private final Map<String, Map<String, SseEmitter>> sseMap = new ConcurrentHashMap<>();
    
    public void addSession(String userId, String clientId, SseEmitter emitter) {
        sseMap.computeIfAbsent(userId, k -> new ConcurrentHashMap<>())
              .put(clientId, emitter);
    }
    
    public void removeSession(String userId, String clientId) {
        if (sseMap.containsKey(userId)) {
            sseMap.get(userId).remove(clientId);
            if (sseMap.get(userId).isEmpty()) {
                sseMap.remove(userId);
            }
        }
    }
    
    public int getClientCount(String connectId) {
        return sseMap.getOrDefault(connectId, Map.of()).size();
    }
    
    public void sendMessageWithUserId(String userId, Object message) {
        sseMap.getOrDefault(userId, Map.of()).forEach((clientId, emitter) -> {
            try {
                emitter.send(message);
            } catch (Exception e) {
                removeSession(userId, clientId);
            }
        });
    }
    public void sendMessageWithClientId(String clientId, Object message) {

        //通过clientId解析出 userId
        String userId = clientId.split("_")[0];

        // 获取用户对应的所有客户端连接
        Map<String, SseEmitter> userClients = sseMap.getOrDefault(userId, Map.of());
        // 获取指定客户端的 SseEmitter
        SseEmitter emitter = userClients.get(clientId);
        if (emitter != null) {
            try {
                // 发送消息
                emitter.send(message);
            } catch (Exception e) {
                // 发送失败时移除连接
                removeSession(userId, clientId);
            }
        }
    }

    public int getTotalConnections() {
        return sseMap.values().stream()
                .mapToInt(Map::size)
                .sum();
    }
}